from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

#语料库 ，每一行一个文档
corpus=['second second document.','second third document.','second second document good.']  

# 构建DataFrame格式数据
corpus = np.array(corpus)
Tf = TfidfVectorizer(use_idf=True)
X = Tf.fit_transform(corpus)

vocs = Tf.get_feature_names()
corpus_array = Tf.transform(corpus).toarray()
#array([[0.50154891, 0.50154891, 0.70490949],
    #[0.4472136 , 0.89442719, 0.        ]])

#corpus_norm_df = pd.DataFrame(corpus_array, columns=vocs)
#print(corpus_norm_df.head())
    #document    second     third
#0  0.501549  0.501549  0.704909
#1  0.447214  0.894427  0.000000

# 使用TfidVectorizer进行TF-idf词袋模型的构建
# 计算相似度
similarity_matrix = cosine_similarity(corpus_array)
similarity_matrix_df = pd.DataFrame(similarity_matrix)
print(similarity_matrix_df)


#要找5个最接近的相关文档
related_docs_indices = similarity_matrix.argsort()[:-5:-1]
print(related_docs_indices)
#[[1 0 2]
#[2 0 1]
#[1 2 0]]